import coremltools
# from sklearn.linear_model import LinearRegression
# import pandas as pd

# coremltools 文档地址：https://apple.github.io/coremltools/
# 转换后的模型以mlmodel为后缀

# Caffe model
# 支持文件类型：bvlc_alexnet.caffemodel, deploy.prototxt, class_labels.txt.
def caffeModel(caffe_path, mlmodel_path):
	coreml_model = coremltools.converters.caffe.convert(caffe_path)
	# modelVisual(coreml_model)
	coremltools.utils.save_spec(coreml_model, mlmodel_path)

	# # Convert a Caffe model to a classifier in Core ML
	# coreml_model = coremltools.converters.caffe.convert(
	# 	('bvlc_alexnet.caffemodel', 'deploy.prototxt'), predicted_feature_name='class_labels.txt'
	# )

	# # Now save the model
	# # 'BVLCObjectClassifier.mlmodel'
	# coreml_model.save(mlmodel_path)

# scikit-learn
def scikitLearn(slearn_path,):
	# Load data
	data = pd.read_csv('houses.csv')

	# Train a model
	model = LinearRegression()
	model.fit(data[["bedroom", "bath", "size"]], data["price"])

	# Convert and save the scikit-learn model
	coreml_model = coremltools.converters.sklearn.convert(model, ["bedroom", "bath", "size"], "price")

# 设置转换模型数据
def metadataSet(model):
	model.author = 'John Smith'
	model.license = 'BSD'
	model.short_description = 'Predicts the price of a house in the Seattle area.'

	# Set feature descriptions manually
	model.input_description['bedroom'] = 'Number of bedrooms'
	model.input_description['bathrooms'] = 'Number of bathrooms'
	model.input_description['size'] = 'Size (in square feet)'

	# Set the output descriptions
	model.output_description['price'] = 'Price of the house'

	# Save the model
	model.save('HousePricer.mlmodel')

# 模型评估验证
def modelEvalution(mlmodel_name):
	# Load the model
	model = coremltools.models.MLModel(mlmodel_name)

	# Make predictions
	predictions = model.predict({'bedroom': 1.0, 'bath': 1.0, 'size': 1240})

# 设置支持可视化
def modelVisual(mlmodel):
	mlmodel.visualize_spec()


if __name__ == '__main__':
	caffeModel('squeezenet_v1.0.caffemodel','squeezenet_v1.0.mlmodel')









